#!/usr/bin/env python
# -*- coding: utf-8 -*-
# @Time    : 2017/11/13 10:51
# @Author  : Yunhao.Cao
# @File    : scheduler.py
from __future__ import absolute_import, unicode_literals
from functools import partial
import multiprocessing
import multiprocessing.pool

__author__ = 'Yunhao.Cao'

__ALL__ = ["Scheduler", ]


class Scheduler:
    def __init__(self, size=None, subprocess=None, sub_kwargs=None, **kwargs):
        """
        初始化进程池
        :param size(int): 进程的数量
        :param subprocess(function(process_id, ...)): 子进程
        :param sub_kwargs(dict): 子进程的参数
        :param kwargs(dict): 参数
        :return:
        """
        if not hasattr(subprocess, "__call__"):
            raise

        if not isinstance(size, (int,)):
            size = 4

        if not isinstance(sub_kwargs, (dict,)):
            sub_kwargs = {}

        # 进程池
        self._size = size
        self._pool = multiprocessing.Pool(self._size)
        self._subprocess = subprocess

        partial_subprocess = partial(subprocess, **sub_kwargs)
        self._pool.map(partial_subprocess, range(self._size))

    def join(self):
        """
        运行进程池
        :return:
        """
        self._pool.close()
        self._pool.join()


def _test(process_id, a):
    print("[{}] {}".format(process_id, a))


if __name__ == '__main__':
    s = Scheduler(subprocess=_test, a=1)
    s.join()
